package com.hu.ta.assignment5.beans;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import com.hu.ta.assignment5.io.FileIO;

public class PosTagger implements Serializable {

	private static final long serialVersionUID = -7399160437769406144L;

	private static final String FINAL_TAG = ".";
	private static final String FIRST_TAG = "$FIRST_TAG$";

	private NGrams transitions = new NGrams();
	private WordPos words = new WordPos();
	private String allTags[];

	public PosTagger() {

	}

	private void addToTransitions(NGram ngram, AnnotatedToken tok) {
		if (transitions.containsKey(ngram)) {
			transitions.get(ngram).add(tok.getPos());
		} else {
			PosProbabilities prob = new PosProbabilities();
			prob.add(tok.getPos());
			transitions.put(ngram, prob);
		}

		if (words.containsKey(tok.getToken())) {
			words.get(tok.getToken()).add(tok.getPos());
		} else {
			PosProbabilities prob = new PosProbabilities();
			prob.add(tok.getPos());
			words.put(tok.getToken(), prob);
		}
	}

	public void train(List<FileIO> trains) {
		List<String> strs = new ArrayList<String>(1);
		strs.add(FIRST_TAG);
		NGram ngram = new NGram(strs);
		for (FileIO f : trains) {
			for (AnnotatedToken tok : f.getText().getTokens()) {
				// Füge neues Wort zu Wahrscheinlichkeiten hinzu.
				addToTransitions(ngram, tok);
				ngram = ngram.getNextNGram(tok.getPos());
				if (tok.isEOS()) {
					ngram = new NGram(strs);
				}
			}
		}

		for (PosProbabilities prob : transitions.values()) {
			prob.frequencies2Probabilities();
		}

		for (PosProbabilities prob : words.values()) {
			prob.frequencies2Probabilities();
		}

		allTags = transitions.getAllTags();
	}

	private String[] getNextSentence(AnnotatedText text) {
		List<String> li = text.getNextSentence();
		String strs[] = new String[li.size()];
		return li.toArray(strs);
	}

	public AnnotatedText tag(AnnotatedText text) {
		AnnotatedText res = new AnnotatedText();

		String[] sentence;
		while ((sentence = getNextSentence(text)).length > 0) {
			AnnotatedText taggedSentence = tagSentence(sentence,
					new PosProbTable(sentence, allTags, words));
			if (null != taggedSentence) {
				res.addText(taggedSentence);
			} else {
				System.out.println("Not tagged sentence.");
			}
		}
		return res;
	}

	private AnnotatedText tagSentence(final String sentence[], PosProbTable ppt) {
		List<String> resTags = new ArrayList<String>();

		final Viterbi<Double> viterbi = new Viterbi<Double>();
		Viterbi<Integer> backpointer = new Viterbi<Integer>();

		boolean bFirst = true;
		Map<String, Double> vi = new HashMap<String, Double>();
		Map<String, Integer> bp = new HashMap<String, Integer>();

		// Initialisierung
		for (String tag : allTags) {
			if (bFirst) {
				viterbi.put(sentence[0], vi);
				backpointer.put(sentence[0], bp);
				bFirst = false;
			}
			vi.put(tag,
					transitions.get(new NGram(FIRST_TAG)).getProbability(tag)
							* words.get(sentence[0]).getProbability(tag));
			bp.put(tag, 0);
		}

		// Rekursion
		for (int w = 1; w < sentence.length; w++) {
			bFirst = true;
			for (final String tag : allTags) {
				// Initialisierung
				if (bFirst) {
					viterbi.put(sentence[w], new HashMap<String, Double>());
					backpointer
							.put(sentence[w], new HashMap<String, Integer>());
					bFirst = false;
				}
				Map<String, Double> vit = viterbi.get(sentence[w]);
				double e = words.get(sentence[w]).getProbability(tag);
				vit.put(tag,
						e
								* getMaxDbl(transitions,
										viterbi.get(sentence[w - 1])));

				Map<String, Integer> backpoint = backpointer.get(sentence[w]);
				backpoint
						.put(tag, getIndexForMax(viterbi.get(sentence[w - 1])));
			}
			bFirst = false;
		}

		viterbi.get(sentence[sentence.length - 1]).put(
				FINAL_TAG,
				getMaxDbl(transitions,
						viterbi.get(sentence[sentence.length - 1])));

		backpointer.get(sentence[sentence.length - 1]).put(FINAL_TAG,
				getIndexForMax(viterbi.get(sentence[sentence.length - 1])));

		// Pfadermittlung
		double maxProb = 0.0;
		int maxProbIdx = -1;
		int i = 0;
		for (String tag : allTags) {
			double prob = 1.0;
			for (int w = sentence.length - 1; w >= 1; w--) {
				prob *= viterbi.get(sentence[w - 1]).get(
						allTags[backpointer.get(sentence[w]).get(tag)]);
			}

			if (maxProb <= prob) {
				maxProbIdx = i;
				maxProb = prob;
			}

			i++;
		}

		if (maxProbIdx > -1) {
			resTags.add(allTags[maxProbIdx]);
			for (int w = sentence.length - 1; w >= 1; w--) {
				resTags.add(
						0,
						allTags[backpointer.get(sentence[w]).get(
								allTags[maxProbIdx])]);
			}
		}

		// Resultat erzeugen
		AnnotatedText taggedText = new AnnotatedText();
		for (int j = 0; j < resTags.size(); j++) {
			String token = sentence[j];
			String tag = resTags.get(j);
			taggedText.addTaggedToken(token, tag);
		}

		return taggedText;
	}

	private static double getMaxDbl(NGrams ngrams, Map<String, Double> tags) {
		double max = 0.0;
		for (String tag : tags.keySet()) {
			max = Math.max(max, ngrams.get(new NGram(FIRST_TAG))
					.getProbability(tag) * tags.get(tag));
		}
		return max;
	}

	private static int getIndexForMax(Map<String, Double> tags) {

		double max = 0.0;
		Map<Double, Integer> valuesMap = new HashMap<Double, Integer>();
		int i = 0;
		for (String tag : tags.keySet()) {
			double d = tags.get(tag);
			valuesMap.put(d, i++);
			max = Math.max(max, d);
		}

		if (null != valuesMap.get(max)) {
			return valuesMap.get(max);
		}

		return -1;
	}
}
